#!/usr/bin/env python

import os
import time
import pylab #@UnresolvedImport
import socket
import urllib2
import httplib2
import urlparse
import platform

from colors import * #@UnusedWildImport
from BeautifulSoup import BeautifulSoup, SoupStrainer

class Shelob():
    
    def __init__(self):
        self.maxPages = 50
        self.timeout = 10
        self.tagDownloads = 1
        self.hrefDownloads = 1
        self.failedDownloads = 1
        self.pages = []
        self.fileExtension = ''
        self.exclude = []
        self.exclusive = []
        self.saveDirectory = None
    
    def loadFile(self):
        if os.path.exists(self.urlFile) and os.path.isfile(self.urlFile):
            file = open(self.urlFile, 'r')
            data = file.readlines()
            for line in data:
                self.pages.append(line.strip())
            file.close()
            print ' [*] Loaded %d Url(s) from file' % len(self.pages)
        else:
            print R+' [!] Error: File (%s) does not exist, see --help' % self.urlFile, W
    
    def addUrl(self, url):
        self.pages.append(url)
        print ' [*] Added %s to spider page(s)' % url
    
    def addExclusion(self, string):
        if not self.isExcluded(string) and len(string) > 0:
            self.exclude.append(string)
            print ' [*] Added "%s" to exclusion list' % string
    
    def addExclusive(self, string):
        if len(string) > 0:
            self.exclusive.append(string)
            print ' [*] Added "%s" to the exclusive list'
    
    def disableColors(self):
        disableColors()
    
    def startCollecting(self, monitor = False):
        ''' Executes the main Spider, starts worker threads '''
        if len(self.pages) == 0:
            print '\n [*] No pages to spider, thanks for trying though!'
            return
        if len(self.fileExtension) == 0:
            print ' \n [*] No file extension set'
            return
        initTime = time.time()
        self.setSaveDirectory()
        self.__executeSpider__(monitor)
        elapsedTime = time.time() - initTime
        itemCount = len(os.listdir(self.saveDirectory))
        print ' [*] Spider completed successfully after %d second(s), and got %d item(s)' % (elapsedTime, itemCount)
        self.pages = []
        self.exclude = []
        if monitor:
            pylab.close()
        
    def setSaveDirectory(self, path = None):
        ''' Set the path where files will be stored '''
        if path == None:
            if platform.system() == 'Linux':
                self.saveDirectory = '/tmp/shelob/'
            else:
                self.setSaveDirectory(os.environ['TEMP']+'/shelob')
        if not os.path.exists(self.saveDirectory):
            os.mkdir(self.saveDirectory)
        print ' [*] Save directory set to:', self.saveDirectory
        
    def isExcluded(self, url):
        ''' Checks to see if a url contains an excluded value '''
        for exclusion in self.exclude:
            if exclusion in url:
                return True
        return False
    
    def __executeSpider__(self, monitor):
        ''' Main spider method, this should only be called from startCollecting '''
        if monitor:
            self.create_monitor()
            self.update_monitor()
        try:
            for page in self.pages:
                self.__collectImages__(page)
                self.__collectLinks__(page)
                print " [*] Completed download(s) for", page, 
                print O+" [%d of %d]" % (self.pages.index(page)+1,len(self.pages)), W
                if monitor:
                    self.update_monitor()
        except KeyboardInterrupt:
            print '\b\b '+R+'[!] Warning: User abort, stopping spider...'+W
            return
        except UnboundLocalError:
            print R+'\n [!] OMFG: Spider crashed, trying to recover'+W
            return
    
    def __collectImages__(self, url):
        print " [*] Collecting images from", url
        try:
            soup = BeautifulSoup(urllib2.urlopen(url, timeout=self.timeout))
            parsed = list(urlparse.urlparse(url))
            self.__downloadImages__(url, soup, parsed)
        except urllib2.HTTPError:
            print R+" [!] Error: Access denied", url+W
            return
        except UnicodeEncodeError:
            print R+" [!] Error: Unable to encode url properly"+W
            return
        except urllib2.URLError:
            print R+' [!] Error: The HTTP request timed out after %d second(s)' % self.timeout, W
            return
        
    def __downloadImages__(self, url, soup, parsed):
        ''' Downloads any embedded images in the page '''
        try:
            for image in soup.findAll("img"):
                try:
                    filename = image["src"].split("/")[-1]
                    parsed[2] = image["src"]
                    outpath = os.path.join(self.saveDirectory, filename)
                    imageUrl = "%(src)s" % image
                    if self.isExcluded(imageUrl):
                        raise TypeError
                    if imageUrl[-3:] == self.fileExtension or self.fileExtension == '*':
                        print " [*] Get image: %s" % imageUrl
                        self.__getFile__(imageUrl, outpath)
                        self.tagDownloads += 1
                except KeyError:
                    print R+' [!] Error: Failed to parse url properly'+W
                    pass
                except TypeError:
                    print G+' [-] Skip: The url contains an excluded string'+W
                    pass
                except ValueError:
                    print R+' [!] Error: Unable to download image, looks like a broken link'+W
                    pass
                except UnicodeEncodeError:
                    print R+' [!] Error: Unable to encode url properly, skipping'+W
                    pass
                except urllib2.URLError:
                    print R+' [!] Error: The http request timed out after %s second(s)' % self.timeout, W
                    self.failedDownloads += 1
                    pass
        except UnboundLocalError:
            print R+" [!] Error: The computer has become self-aware, RUN AWAY!"+W
            pass
    
    def __collectLinks__(self, url):
        ''' Collect all the links on a page '''
        try:
            http = httplib2.Http()
            print " [*] Checking %s for links" % url
            status, response = http.request(url) #@UnusedVariable
            for link in BeautifulSoup(response, parseOnlyThese = SoupStrainer('a')):
                if link.has_key('href') and link['href'].startswith("http://"):
                    self.__handleHref__(link)
            print " [*] Completed link check for", url
        except httplib2.RedirectLimit:
            print R+' [!] Error: Too many redirects from', url+W
            pass
    
    def __handleHref__(self, link):
        ''' Downloads any images linked to using href= '''
        if link['href'].endswith(self.fileExtension):
            try:
                fileName = link['href'].split("/")[-1]
                if not os.path.exists(self.saveDirectory + fileName):
                    if not self.isExcluded(link['href']):
                        print ' [*] Get href:', link['href']
                        self.__getFile__(link['href'], self.saveDirectory + fileName)
                        self.hrefDownloads += 1
            except:
                print R+' [!] Error: Unable to download file from href', link['href']+W
                self.failedDownloads += 1
                pass
        elif not link['href'] in self.pages and len(self.pages) < self.maxPages:
            self.pages.append(link['href'])
            print " [*] Added new link:", link['href']
    
    def __getFile__(self, url, saveTo):
        ''' Download a remote file from url and writes it to "saveTo" '''
        try:
            headers = {'User-Agent' : 'Mozilla/5.0 (compatible; MSIE 9.0; Windows NT 6.0; Trident/5.0'}
            req = urllib2.Request(url, None, headers)
            remoteSite = urllib2.urlopen(req, timeout = self.timeout)
            file = open(saveTo, 'w')
            while True:
                buffer = remoteSite.read(2048)
                if not buffer:
                    break
                file.write(buffer)
            file.close()
        except socket.timeout:
            os.remove(saveTo)
            print R+' [!] Warning: Download failed/interrupted, removed partial file'+W
            pass
    
    def create_monitor(self):
        ''' Creates a  realtime progress monitor '''
        pylab.ion()           # Starts interactive mode
        fig = pylab.figure(1) #@UnusedVariable
        self.chart = pylab.subplot(111)
        ind = pylab.arange(1,4)
        self.tagBar, self.hrefBar, self.failedBar = pylab.bar(ind, self.tagDownloads, self.hrefDownloads, self.failedDownloads)
        centers = ind + 0.5*self.tagBar.get_width()
        self.tagBar.set_facecolor('g')
        self.hrefBar.set_facecolor('b')
        self.failedBar.set_facecolor('r')
        self.chart.set_xlim([0.5,4])
        self.chart.set_xticks(centers)
        self.chart.set_ylim([0,100])
        self.chart.set_xticklabels(['Embedded Downloads', 'Href Downloads', 'Failed Downloads'])
        self.chart.set_ylabel('Number of Images')
        self.chart.set_title('Shelob Realtime Monitor')
       
    def update_monitor(self):
        ''' Updates the monitor with the most recent stats '''
        print ' [*] Status update sent to monitor'
        self.tagBar.set_height(self.tagDownloads)
        self.hrefBar.set_height(self.hrefDownloads)
        self.failedBar.set_height(self.failedDownloads)
        if self.hrefDownloads + self.tagDownloads > 100:
            total = self.hrefDownloads + self.tagDownloads + 10
            self.chart.set_ylim([0,total])
        else:
            self.chart.set_ylim([0,100])
        pylab.draw()

